package handler

import (
	"fmt"
	"github.com/gin-gonic/gin"
	"gorm.io/gorm"
	"myai/core"
	"myai/core/types"
	"myai/service"
	"myai/service/mj"
	"myai/service/oss"
	"myai/store/model"
	"myai/store/model/req"
	"myai/store/vo"
	"myai/utils"
	"myai/utils/resp"
	"strings"
	"time"
)

type MidJourneyHandler struct {
	BaseHandler
	mjService   *mj.Service
	snowflake   *service.Snowflake
	uploader    *oss.UploaderManager
	userService *service.UserService
}

func NewMidJourneyHandler(app *core.AppServer, db *gorm.DB, snowFlake *service.Snowflake, service *mj.Service,
	manager *oss.UploaderManager, userService *service.UserService) *MidJourneyHandler {
	return &MidJourneyHandler{
		snowflake:   snowFlake,
		mjService:   service,
		uploader:    manager,
		userService: userService,
		BaseHandler: BaseHandler{
			App: app,
			DB:  db,
		},
	}
}

// preCheck 前置检查
func (h *MidJourneyHandler) preCheck(c *gin.Context) bool {
	user, err := h.GetLoginUser(c)
	if err != nil {
		resp.NotAuth(c)
		return false
	}
	if user.Power < h.App.SysConfig.MjPower {
		resp.ERROR(c, "当前用户剩余算力不足以完成本次绘画！")
		return false
	}
	return true
}

// Image godoc
// @Summary      创建一个绘画任务
// @Description  创建一个绘画任务
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        req   body     req.MjImageReq  true  "绘画参数"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/mj/image [post]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) Image(c *gin.Context) {
	var data req.MjImageReq
	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	if !h.preCheck(c) {
		return
	}

	var params = ""
	if data.Rate != "" && !strings.Contains(params, "--ar") {
		params += " --ar " + data.Rate
	}

	if data.Seed > 0 && !strings.Contains(params, "--seed") {
		params += fmt.Sprintf(" --seed %d", data.Seed)
	}

	if data.Stylize > 0 && !strings.Contains(params, "--s") && !strings.Contains(params, "--stylize") {
		params += fmt.Sprintf(" --s %d", data.Stylize)
	}

	if data.Chaos > 0 && !strings.Contains(params, "--c") && !strings.Contains(params, "--chaos") {
		params += fmt.Sprintf(" --c %d", data.Chaos)
	}

	if len(data.ImgArr) > 0 && data.Iw > 0 {
		params += fmt.Sprintf(" --iw %.2f", data.Iw)
	}

	if data.Raw {
		params += " --style raw"
	}

	if data.Quality > 0 {
		params += fmt.Sprintf(" --q %.2f", data.Quality)
	}

	if data.Tile {
		params += " --tile"
	}

	if data.CRef != "" {
		params += fmt.Sprintf(" --cref %s", data.CRef)
		if data.Cw > 0 {
			params += fmt.Sprintf(" --cw %d", data.Cw)
		} else {
			params += " --cw 100"
		}
	}

	if data.SRef != "" {
		params += fmt.Sprintf(" --sref %s", data.SRef)
	}

	// --v 使用的模型  --niji 使用niji动漫风格模型
	if data.Model != "" && !strings.Contains(params, "--v") && !strings.Contains(params, "--niji") {
		params += fmt.Sprintf(" %s", data.Model)
	}

	// 处理融图和换脸的提示词
	if data.TaskType == types.TaskSwapFace.String() || data.TaskType == types.TaskBlend.String() {
		params = fmt.Sprintf("%s:%s", data.TaskType, strings.Join(data.ImgArr, ","))
	}

	// 如果本地图片上传的是相对路径，处理成绝对路径
	for i, v := range data.ImgArr {
		if !strings.HasPrefix(v, "http") {
			data.ImgArr[i] = fmt.Sprintf("http://localhost:5679/%s", strings.TrimLeft(v, "/"))
		}
	}

	idValue, _ := c.Get(types.LoginUserID)
	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
	taskId, err := h.snowflake.Next(true)
	if err != nil {
		resp.ERROR(c, "生成任务Id失败："+err.Error())
		return
	}

	task := types.MjTask{
		ClientId:         data.ClientId,
		TaskId:           taskId,
		Type:             types.TaskType(data.TaskType),
		Prompt:           data.Prompt,
		NegPrompt:        data.NegPrompt,
		Params:           params,
		UserId:           userId,
		ImgArr:           data.ImgArr,
		Mode:             h.App.SysConfig.MjMode,
		TranslateModelId: h.App.SysConfig.TranslateModelId,
	}

	job := model.MidJourneyJob{
		Type:      data.TaskType,
		UserId:    uint(userId),
		TaskId:    taskId,
		TaskInfo:  utils.JsonEncode(task),
		Progress:  0,
		Prompt:    fmt.Sprintf("%s %s", data.Prompt, params),
		Power:     h.App.SysConfig.MjPower,
		CreatedAt: time.Now(),
	}

	opt := "绘图"
	if data.TaskType == types.TaskBlend.String() {
		job.Prompt = "融图：" + strings.Join(data.ImgArr, ",")
		opt = "融图"
	} else if data.TaskType == types.TaskSwapFace.String() {
		job.Prompt = "换脸：" + strings.Join(data.ImgArr, ",")
		opt = "换脸"
	}

	if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
		resp.ERROR(c, "添加任务失败："+res.Error.Error())
		return
	}

	task.Id = job.Id
	// 推送任务
	h.mjService.PushTask(task)

	// 扣减算力
	err = h.userService.DecreasePower(int(job.UserId), job.Power, model.PowerLog{
		Type:   types.PowerConsume,
		Model:  "mid-journey",
		Remark: fmt.Sprintf("%s操作，任务ID：%s", opt, job.TaskId),
	})
	if err != nil {
		resp.ERROR(c, "扣减用户算力失败："+err.Error())
		return
	}
	resp.SUCCESS(c)
}

// Upscale godoc
// @Summary      向MidJourney Bot发送高级命令
// @Description  向MidJourney Bot发送高级命令
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        req   body     req.MjReq  true  "绘画指令"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/mj/upscale [post]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) Upscale(c *gin.Context) {
	var data req.MjReq
	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	if !h.preCheck(c) {
		return
	}

	idValue, _ := c.Get(types.LoginUserID)
	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
	taskId, _ := h.snowflake.Next(true)
	task := types.MjTask{
		Type:        types.TaskUpscale,
		ClientId:    data.ClientId,
		UserId:      userId,
		Index:       data.Index,
		ChannelId:   data.ChannelId,
		MessageId:   data.MessageId,
		MessageHash: data.MessageHash,
		Mode:        h.App.SysConfig.MjMode,
	}

	job := model.MidJourneyJob{
		Type:      types.TaskUpscale.String(),
		ChannelId: data.ChannelId,
		UserId:    uint(userId),
		TaskId:    taskId,
		TaskInfo:  utils.JsonEncode(task),
		Progress:  0,
		Power:     h.App.SysConfig.MjActionPower,
		CreatedAt: time.Now(),
	}

	if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
		resp.ERROR(c, "添加任务失败："+res.Error.Error())
		return
	}

	task.Id = job.Id
	h.mjService.PushTask(task)
	err := h.userService.DecreasePower(int(job.UserId), job.Power, model.PowerLog{
		Type:   types.PowerConsume,
		Model:  "mid-journey",
		Remark: fmt.Sprintf("Upscale 操作，任务ID：%s", job.TaskId),
	})
	if err != nil {
		resp.ERROR(c, "扣减算力失败："+err.Error())
		return
	}
	resp.SUCCESS(c)
}

// Variation godoc
// @Summary      向MidJourney Bot发送变异命令
// @Description  向MidJourney Bot发送变异命令
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        req   body     req.MjReq  true  "变异绘画指令"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/mj/variation [post]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) Variation(c *gin.Context) {
	var data req.MjReq
	if err := c.ShouldBindJSON(&data); err != nil {
		resp.ERROR(c, types.InvalidArgs)
		return
	}

	if !h.preCheck(c) {
		return
	}

	idValue, _ := c.Get(types.LoginUserID)
	userId := utils.IntValue(utils.InterfaceToString(idValue), 0)
	taskId, _ := h.snowflake.Next(true)
	task := types.MjTask{
		Type:        types.TaskVariation,
		ClientId:    data.ClientId,
		UserId:      userId,
		Index:       data.Index,
		ChannelId:   data.ChannelId,
		MessageId:   data.MessageId,
		MessageHash: data.MessageHash,
		Mode:        h.App.SysConfig.MjMode,
	}

	job := model.MidJourneyJob{
		Type:      types.TaskVariation.String(),
		ChannelId: data.ChannelId,
		UserId:    uint(userId),
		TaskId:    taskId,
		TaskInfo:  utils.JsonEncode(task),
		Progress:  0,
		Power:     h.App.SysConfig.MjActionPower,
		CreatedAt: time.Now(),
	}

	if res := h.DB.Create(&job); res.Error != nil || res.RowsAffected == 0 {
		resp.ERROR(c, "添加任务失败："+res.Error.Error())
		return
	}

	task.Id = job.Id
	h.mjService.PushTask(task)
	err := h.userService.DecreasePower(int(job.UserId), job.Power, model.PowerLog{
		Type:   types.PowerConsume,
		Model:  "mid-journey",
		Remark: fmt.Sprintf("Variation 操作，任务ID：%s", job.TaskId),
	})
	if err != nil {
		resp.ERROR(c, "扣减算力失败："+err.Error())
		return
	}
	resp.SUCCESS(c)
}

// ImgWall godoc
// @Summary      照片墙
// @Description  照片墙
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        page   query    int  true  "页码"
// @Param        page_size   query    int  true  "每页数量"
// @Success      200  {object}  types.BizVo{data=vo.Page}
// @Router       /api/mj/imgWall [get]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) ImgWall(c *gin.Context) {
	page := h.GetInt(c, "page", 1)
	pageSize := h.GetInt(c, "page_size", 10)
	err, jobs := h.getData(true, 0, page, pageSize, true)
	if err != nil {
		resp.ERROR(c, "获取照片墙数据失败："+err.Error())
		return
	}
	resp.SUCCESS(c, jobs)
}

// JobList godoc
// @Summary      获取MJ任务列表
// @Description  获取MJ任务列表
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        page        query    int   true  "页码"
// @Param        page_size   query    int   true  "每页数量"
// @Param        finish      query    bool  true  "任务是否完成"
// @Param        publish     query    bool  true  "是否发布"
// @Success      200  {object}  types.BizVo{data=vo.Page}
// @Router       /api/mj/jobs [get]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) JobList(c *gin.Context) {
	finish := h.GetBool(c, "finish")
	userId := h.GetLoginUserId(c)
	page := h.GetInt(c, "page", 0)
	pageSize := h.GetInt(c, "page_size", 0)
	publish := h.GetBool(c, "publish")

	err, jobs := h.getData(finish, userId, page, pageSize, publish)
	if err != nil {
		resp.ERROR(c, "获取列表数据失败："+err.Error())
		return
	}
	resp.SUCCESS(c, jobs)
}

// Remove godoc
// @Summary      删除任务
// @Description  删除任务
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        id        query    int   true  "任务Id"
// @Param        user_id   query    int   true  "用户Id"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/mj/remove [get]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) Remove(c *gin.Context) {
	id := h.GetInt(c, "id", 0)
	userId := h.GetInt(c, "user_id", 0)
	var job model.MidJourneyJob
	if res := h.DB.Where("id = ? and user_id = ?", id, userId).First(&job); res.Error != nil {
		resp.ERROR(c, "记录不存在")
		return
	}

	// 删除任务
	err := h.DB.Delete(&job).Error
	if err != nil {
		resp.ERROR(c, "删除任务出现错误："+err.Error())
		return
	}

	// 删除图片
	err = h.uploader.GetUploadHandler().Delete(job.ImgURL)
	if err != nil {
		log.Errorf("删除图片出现错误：%v", err)
	}
	resp.SUCCESS(c)
}

// Publish godoc
// @Summary      发布图片到画廊显示
// @Description  发布图片到画廊显示
// @Tags         前台-MJ绘画
// @Accept       json
// @Produce      json
// @Param        id          query    int   true  "任务Id"
// @Param        user_id     query    int   true  "用户Id"
// @Param        action      query    bool  true  "是否发布"
// @Success      200  {object}  types.BizVo{}
// @Router       /api/mj/publish [get]
// @Security ApiKeyAuth
func (h *MidJourneyHandler) Publish(c *gin.Context) {
	id := h.GetInt(c, "id", 0)
	userId := h.GetInt(c, "user_id", 0)
	// 发布动作，true => 发布，false => 取消分享
	action := h.GetBool(c, "action")
	err := h.DB.Model(&model.MidJourneyJob{Id: uint(id), UserId: uint(userId)}).UpdateColumn("publish", action).Error
	if err != nil {
		resp.ERROR(c, "操作数据库出现错误："+err.Error())
		return
	}
	resp.SUCCESS(c)
}

// getData 获取 MJ 任务列表
func (h *MidJourneyHandler) getData(finish bool, userId uint, page int, pageSize int, publish bool) (error, vo.Page) {
	session := h.DB.Session(&gorm.Session{})
	if finish {
		session = session.Where("progress >= ?", 100).Order("id desc")
	} else {
		session = session.Where("progress < ?", 100).Order("id asc")
	}

	if userId > 0 {
		session = session.Where("user_id = ?", userId)
	}

	if publish {
		session = session.Where("publish = ?", publish)
	}

	if page > 0 && pageSize > 0 {
		offset := (page - 1) * pageSize
		session = session.Offset(offset).Limit(pageSize)
	}

	// 统计总数
	var total int64
	session.Model(&model.MidJourneyJob{}).Count(&total)

	var items []model.MidJourneyJob
	res := session.Find(&items)
	if res.Error != nil {
		return res.Error, vo.Page{}
	}

	var jobs = make([]vo.MidJourneyJob, 0)
	for _, item := range items {
		var job vo.MidJourneyJob
		err := utils.CopyObject(item, &job)
		if err != nil {
			log.Errorf("拷贝数据出现错误：%v", err.Error())
			continue
		}
		jobs = append(jobs, job)
	}
	return nil, vo.NewPage(total, page, pageSize, jobs)
}
